#pragma once

#include <algorithm/analytical_model/Constants.h>
#include <algorithm/analytical_model/algorithm/cuda/inter_op/Data.h>

#include <memory>

namespace solar
{
    namespace cuda
    {
        class ConvolutionMethod;
    }
    namespace dnn
    {
        struct UNIZAR2Output;
    }
} // namespace solar

namespace solar::cuda::inter_op
{
    class ConvolutionMethod
    {
        public:
            // 圆柱形接收器 + 矩形定日镜
            ConvolutionMethod(Float3* helio_verts, Float3 helio_normal, Float3 helio_pos,
                              Float3 recv_pos, Float3 helio_size, Float3 recv_size,
                              Float3 focus_center, Int2 helio_row_col, Int2 recv_row_col,
                              bool* d_micro_sb = nullptr,
                              ConvolutionModelType mode = ConvolutionModelType::kUNIZAR,
                              float* d_array = nullptr, cudaStream_t stream = nullptr);

            // 矩形接收器 + 矩形定日镜
            ConvolutionMethod(Float3* helio_verts, Float3 helio_normal, Float3 helio_position,
                              Float3 focus_center, Int2 helio_row_col, Int2 recv_row_col,
                              Float3* recv_verts, Float3 recv_normal, bool* d_micro_sb = nullptr,
                              float* d_array = nullptr, cudaStream_t stream = nullptr);

            /// @brief 圆柱形接收器 + 矩形定日镜（整个镜场）
            /// @param d_image_plane_xyz_axes 每个定日镜的成像平面的xyz轴
            /// @param d_helio_verts 整个镜场的定日镜顶点的世界坐标系坐标（每个定日镜4个顶点）
            /// @param d_micro_sb 阴影遮挡计算结果（BitMask）
            /// @param d_flux_params 整个镜场的定日镜的参数（sigma1, sigma2, peak）
            /// @param d_focus_centers 定日镜聚焦中心
            /// @param num_helios 定日镜个数
            /// @param recv_size 圆柱形接收器尺寸（半径，高度，ignore）
            /// @param recv_pos_center 圆柱形接收器中心的世界坐标系坐标
            /// @param helio_row_col 定日镜划分（像素）
            /// @param recv_row_col 接收器划分（像素）
            /// @param mode 卷积模型的类型
            /// @param d_array 结果数组
            /// @param stream CUDA Stream
            ConvolutionMethod(Float3* d_image_plane_xyz_axes, Float3* d_helio_verts,
                              bool* d_micro_sb, bool* d_micro_sb_per_row, Float3* d_flux_params,
                              Float3* d_focus_centers, std::size_t num_helios, Float3 recv_size,
                              Float3 recv_pos_center, Int2 helio_row_col, Int2 recv_row_col,
                              ConvolutionModelType mode = ConvolutionModelType::kUNIZAR,
                              float* d_array = nullptr, cudaStream_t stream = nullptr);

            /// @brief 圆柱形接收器 + 矩形定日镜（整个镜场）
            /// @param d_image_plane_xyz_axes 每个定日镜的成像平面的xyz轴
            /// @param d_helio_verts 整个镜场的定日镜顶点的世界坐标系坐标（每个定日镜4个顶点）
            /// @param d_micro_sb 阴影遮挡计算结果（BitMask）
            /// @param d_model_outputs 神经网络输出数组
            /// @param d_focus_centers 定日镜聚焦中心
            /// @param num_helios 定日镜个数
            /// @param recv_size 圆柱形接收器尺寸（半径，高度，ignore）
            /// @param recv_pos_center 圆柱形接收器中心的世界坐标系坐标
            /// @param helio_row_col 定日镜划分（像素）
            /// @param recv_row_col 接收器划分（像素）
            /// @param mode 卷积模型的类型
            /// @param d_array 结果数组
            /// @param stream CUDA Stream
            ConvolutionMethod(Float3* d_image_plane_xyz_axes, Float3* d_helio_verts,
                              bool* d_micro_sb, bool* d_micro_sb_per_row,
                              ::solar::dnn::UNIZAR2Output* d_model_outputs, Float3* d_focus_centers,
                              std::size_t num_helios, Float3 recv_size, Float3 recv_pos_center,
                              Int2 helio_row_col, Int2 recv_row_col,
                              ConvolutionModelType mode = ConvolutionModelType::kUNIZAR,
                              float* d_array = nullptr, cudaStream_t stream = nullptr);

            // For convenient use only
            ConvolutionMethod() = default;

            template <int N = 10> [[nodiscard]] auto loadGlParam() const -> bool
            {
                constexpr bool available_n = N == 10;
                if constexpr (available_n)
                {
                    return this->loadGlParamImpl(N);
                }
                else
                {
                    static_assert(!available_n, "n must be one of the values: 10");
                }
            }

            /// @brief 绘制圆柱形接收器上的通量图
            void drawFluxMapCylinder() const;
            /// @brief 绘制圆柱形接收器上的通量图
            void drawFluxMapCylinderFast() const;
            /// @brief 绘制平面接收器上的通量图
            void drawFluxMap() const;
            /// @brief 绘制平面接收器上的通量图
            void drawFluxMapFast() const;
            /// @brief 获取仿真结果（数据在GPU上）
            [[nodiscard]] auto getFluxMap() const -> float*;
            [[nodiscard]] auto getCudaStream() const -> cudaStream_t;
            void setCudaStream(const cudaStream_t& stream) const;
            void setModelParams(float sigma1, float sigma2, float peak) const;
            void setModelParams(const dnn::UNIZAR2Output& params) const;
            void setHeliostatID(std::uint64_t helio_id) const;

        private:
            std::shared_ptr<::solar::cuda::ConvolutionMethod> delegate_ptr_;

            [[nodiscard]] auto loadGlParamImpl(int n = 10) const -> bool;
    };
} // namespace solar::cuda::inter_op
